package com.wl.sharding;

import com.google.common.collect.Range;
import org.apache.shardingsphere.sharding.api.sharding.standard.PreciseShardingValue;
import org.apache.shardingsphere.sharding.api.sharding.standard.RangeShardingValue;
import org.apache.shardingsphere.sharding.api.sharding.standard.StandardShardingAlgorithm;

import java.text.SimpleDateFormat;
import java.util.*;

/**
 * @author: zhouyu
 * @date: 2025/4/5
 * @Description:
 */
public class MonthRangeShardingAlgorithm implements StandardShardingAlgorithm {
    @Override
    public String doSharding(Collection availableTargetNames, PreciseShardingValue shardingValue) {
        // 精确分片逻辑（单月查询）
        SimpleDateFormat sdf = new SimpleDateFormat("yyyyMM");
        String tableSuffix = sdf.format(shardingValue.getValue());
        String tableName = shardingValue.getLogicTableName() + "_" + tableSuffix;

        // 检查表是否存在
        if (availableTargetNames.contains(tableName)) {
            return tableName;
        }
        throw new IllegalArgumentException("无效的表名: " + tableName);
    }

    @Override
    public Collection<String> doSharding(Collection availableTargetNames, RangeShardingValue shardingValue) {
        // 范围查询分片逻辑
        Set<String> result = new LinkedHashSet<>();

        // 1. 获取查询范围
        Range<Date> range = shardingValue.getValueRange();
        Date lower = range.hasLowerBound() ? range.lowerEndpoint() : null;
        Date upper = range.hasUpperBound() ? range.upperEndpoint() : null;

        // 2. 如果没有范围（全表查询），返回所有可用表（或根据业务限制）
        if (lower == null && upper == null) {
            return availableTargetNames;
        }

        // 3. 计算涉及的月份范围
        Calendar cal = Calendar.getInstance();
        if (lower != null) {
            cal.setTime(lower);
        } else if (upper != null) {
            cal.setTime(upper);
        }

        // 4. 如果是单月查询（不跨月）
        if (isSameMonth(lower, upper)) {
            String tableName = shardingValue.getLogicTableName() + "_" +
                    new SimpleDateFormat("yyyyMM").format(lower != null ? lower : upper);
            if (availableTargetNames.contains(tableName)) {
                result.add(tableName);
            }
            return result;
        }

        // 5. 处理跨月查询
        Calendar current = (Calendar) cal.clone();
        while (upper == null || !current.getTime().after(upper)) {
            String tableName = shardingValue.getLogicTableName() + "_" +
                    new SimpleDateFormat("yyyyMM").format(current.getTime());
            if (availableTargetNames.contains(tableName)) {
                result.add(tableName);
            }

            // 移动到下个月
            current.add(Calendar.MONTH, 1);
            if (upper != null && current.getTime().after(upper)) {
                break;
            }
        }

        return result;
    }

    private boolean isSameMonth(Date d1, Date d2) {
        if (d1 == null || d2 == null) return false;
        Calendar cal1 = Calendar.getInstance();
        cal1.setTime(d1);
        Calendar cal2 = Calendar.getInstance();
        cal2.setTime(d2);
        return cal1.get(Calendar.YEAR) == cal2.get(Calendar.YEAR) &&
                cal1.get(Calendar.MONTH) == cal2.get(Calendar.MONTH);
    }

    @Override
    public Properties getProps() {
        return null;
    }

    @Override
    public void init(Properties properties) {

    }
}